# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import asyncio
import time
from events import scheduler
from tortoise.transactions import in_transaction
from common.enums.agent import AgentEnum
from common.enums.models import ModelsEnum
from common.enums.wallet import WalletEnum
from common.enums.public import CrontabEnum
from common.cache.queue_cache import QueueCache
from common.models.users import UserModel
from common.models.users import UserWalletModel
from common.models.dev import DevModelsModel
from common.models.sys import SysCrontabModel
from common.models.agents import AgentKnowModel
from common.postgres.public import PgKnowledgeModel
from common.postgres.public import PgAttachmentModel
from common.chain.vector_server import VectorService
from common.chain.chain_server import ChatUtils


async def execute(**kwargs):
    """ 知识库向量训练 """
    start_time = time.time()
    crontab_id: int = int(kwargs["w_id"])
    crontab_job: str = str(kwargs["w_job"])
    crontab_ix: int = int(kwargs["w_ix"])
    if not scheduler.get_job(crontab_job):
        scheduler.pause_job(crontab_job)

    try:
        s = float(f"{(0.001 * crontab_ix):.3f}")
        await asyncio.sleep(s)

        await _handle_embedding()
        await SysCrontabModel.compute(crontab_id, start_time)
    except Exception as e:
        await SysCrontabModel.compute(crontab_id, start_time, status=CrontabEnum.CRON_ERROR, error=str(e))


async def _handle_embedding():
    _uuid = await QueueCache.pop_queue_value(QueueCache.EM_JOB)
    if _uuid is None:
        return True

    # 查询数据
    data = await PgKnowledgeModel.filter(uuid=_uuid, is_delete=0).first()
    if not data:
        return True
    if not data.question:
        return await update_data_error(data, error="学习内容是空的")

    # 查询知识库
    know = await AgentKnowModel.filter(id=data.know_id, is_delete=0).first()
    if not know:
        return await update_data_error(data, error="知识库已被删除")
    if know.is_disable:
        return await update_data_error(data, error="知识库已被禁用")

    # 验证模型
    try:
        models = await DevModelsModel.check_models(know.vector_model, ModelsEnum.TYPE_VECTOR)
    except Exception as e:
        return await update_data_error(data, error=str(e))

    # 查询用户
    user = await UserModel.filter(id=know.user_id, is_delete=0).first()
    if not user:
        return await update_data_error(data, error="账号已注销")
    if user.is_disable:
        return await update_data_error(data, error="账号已冻结")
    if user.points <= 0:
        return await update_data_error(data, error="账号余额不足")

    try:
        # 标记开始
        data.status = AgentEnum.EMB_ING
        data.update_time = int(time.time())
        await data.save()

        # 读取附件
        attr_texts = []
        attachment = await PgAttachmentModel.filter(dataset_id=_uuid, status=AgentEnum.EMB_WAIT).all()
        for att in attachment:
            attr_texts.append(att.question)

        # 内容向量化
        vectorService = VectorService()
        embedding_arr = await vectorService.to_embed(know.vector_model, data.question)
        embedding_str = "[" + ",".join(str(item) for item in embedding_arr) + "]"

        # 附件向量化
        emb_lists = []
        if attr_texts:
            emb_lists = await vectorService.to_embed(know.vector_model, attr_texts)

        # 计算tokens
        usages: dict = vectorService.usage
        use_tokens: int = usages["total_tokens"]
        use_points: int = ChatUtils.compute_price(usages["total_tokens"], models["price"])

        async with in_transaction("pgsql"):
            # 更新PG附件
            for index, item in enumerate(attachment):
                texts: str = attr_texts[index]
                signal_use_tokens: int = ChatUtils.compute_tokens(texts)
                signal_use_points: int = ChatUtils.compute_price(signal_use_tokens, models["price"])
                use_points += signal_use_points
                use_tokens += signal_use_tokens
                emb_arr: list = emb_lists[index]
                emb_str: str = "[" + ",".join(str(s) for s in emb_arr) + "]"
                item.status = AgentEnum.EMB_YES
                item.embedding = emb_str
                item.dimension = len(emb_arr)
                item.vector_model = know.vector_model
                item.vector_alias = models["alias"]
                item.update_time = int(time.time())
                item.use_tokens = signal_use_tokens
                item.use_points = signal_use_points
                await item.save()

            # 更新PG表
            await PgKnowledgeModel.filter(uuid=_uuid).update(
                dimension=len(embedding_arr),
                embedding=embedding_str,
                vector_model=know.vector_model,
                vector_alias=models["alias"],
                use_points=use_points,
                use_tokens=use_tokens,
                status=AgentEnum.EMB_YES,
                update_time=int(time.time())
            )

            # 扣费逻辑
            if use_points:
                # 用户扣费
                user = await UserModel.filter(id=user.id).first()
                surplus_balance: int = max(user.points - use_points, 0)
                await UserModel.filter(id=user.id).update(points=surplus_balance)

                # 写入日志
                await UserWalletModel.dec(
                    scene=WalletEnum.POINTS,
                    user_id=user.id,
                    change_type=WalletEnum.POINTS_DEC_AGENT_EMB,
                    change_amount=use_points,
                    remarks=WalletEnum.get_change_type_msg(WalletEnum.POINTS_DEC_AGENT_EMB),
                    project=know.name,
                    source_sn=know.code,
                    source_id=know.id
                )
    except Exception as e:
        if e.args[0] in [401, 402, 403]:
            await update_data_error(data, error=e.args[1])

        await PgKnowledgeModel.filter(uuid=_uuid).update(
            embedding=None,
            error=e.args[1],
            vector_model=know.vector_model,
            vector_alias=models["alias"],
            status=AgentEnum.EMB_FAIL,
            update_time=int(time.time())
        )


async def update_data_error(data, error: str):
    return await PgKnowledgeModel.filter(uuid=data.uuid).update(
        error=error,
        status=AgentEnum.EMB_FAIL,
        update_time=int(time.time())
    )
